# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Leap-frog integrator
"""

from typing import Tuple

import torch
import math
from torch import Tensor

from .integrator import Integrator, _integrator_register
from ..thermostat import Thermostat
from ..barostat import Barostat
from ..constraint import Constraint
from ...system import Molecule
from ...function import get_arguments, get_tensor


@_integrator_register('nmpimd')
class NMPIMD(Integrator):
    r"""A leap-frog integrator based on "middle scheme" developed by Jian Liu, et al. It is a subclass of `Integrator`.

    Reference:

        Zhang, Z.; Yan, K; Liu, X.; Liu, J.
        A Leap-Frog Algorithm-based Efficient Unified Thermostat Scheme for Molecular Dynamics [J].
        Chinese Science Bulletin, 2018, 63(33): 3467-3483.

    Args:

        system (Molecule):          Simulation system

        thermostat (Thermostat):    Thermostat for temperature coupling. Default: None

        barostat (Barostat):        Barostat for pressure coupling. Default: None

        constraint (Constraint):    Constraint algorithm. Default: None


    Supported Platforms:

        ``Ascend`` ``GPU``

    """
    def __init__(self,
                 system: Molecule,
                 time_step = 1e-3,
                 temperature = 300.,
                 gamma0 = 10.,
                 pressure = 1,
                 tau = 1,
                 pressure_control_step = 1,
                 compressibility = 4.6e-5,
                 remove_com = True,
                 **kwargs
                 ):

        super().__init__(
            system=system,
            time_step=time_step,
            thermostat=None,
            barostat=None,
            constraint=None,
        )

        self.name = 'PILE'
        n = self.num_walker
        self.time_step = time_step
        self.sqrtP = torch.sqrt(get_tensor(n, torch.float32))
        def build_matrix(P):  
            A = get_tensor(torch.zeros(P, P), torch.float32)  # 初始化一个P x P的零矩阵  
            for i in range(P):  
                for j in range(P):  
                    # 使用模P运算来处理边界情况  
                    A[i, j] = 2 * (i == j) - (i  == (j - 1)% P) - (i == (j + 1) % P)  
            return A
        A = build_matrix(n)
        self.eigval, self.eigvec = torch.linalg.eigh(A)
        self.eigval[0] = 0
        self.eigval *= n

        self.temperature = self._get_mw_tensor(temperature, 'temperature')
        boltzmann = system.units.boltzmann
        hbar = 1.054572e-37 * 6.02e23 * 1e12
        self.beta = 1 / (boltzmann * self.temperature)
        self.omega_p = get_tensor(torch.ones(self.num_walker,1,1), torch.float32) * self.sqrtP / (self.beta * hbar)
        self.omega_p[0] = 0
        self.mass_k = (self.atom_mass.repeat(n,1) * self.eigval.unsqueeze(-1)).unsqueeze(-1)
        self.mass_prime = torch.cat([self.atom_mass.unsqueeze(-1), self.mass_k[1:]], dim=0)

        self.gamma0 = gamma0
        self.gamma = torch.cat([torch.tensor([[[self.gamma0]]],device=self.device), self.omega_p[1:]], dim=0)
        self.c1 = torch.exp(-self.gamma*self.time_step)
        self.c2 = torch.sqrt(1-self.c1**2) / self.kinetic_unit_scale
        self.sqrt_mass= torch.sqrt(self.mass_prime)
        self.remove_com = remove_com

        self.pressure = self._get_mw_tensor(pressure, 'pressure')
        self.tau = get_tensor(tau, torch.float32)
        self.compressibility = get_tensor(compressibility, torch.float32)
        self.h_mass_inverse_0 = torch.square(self.tau) / self.compressibility
        self.pressure_control_step = pressure_control_step        

    def forward(self,
                  coordinate: Tensor,
                  velocity: Tensor,
                  force: Tensor,
                  energy: Tensor,
                  kinetics: Tensor,
                  virial: Tensor = None,
                  pbc_box: Tensor = None,
                  step: int = 0,
                  ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:

        # (B,A,D) = (B,A,D) * (B,A,1)

        raise NotImplementedError
    
        # Not checked yet
        C_jk = self.eigvec.to(coordinate.device)
        momentum = velocity * self._atom_mass
        frc = 1/self.sqrtP*torch.einsum("kj,jad->kad",C_jk.T,force)
        p_tilde = 1/self.sqrtP*torch.einsum("jk,jad->kad",C_jk,momentum)
        x_tilde = 1/self.sqrtP*torch.einsum("jk,jad->kad",C_jk,coordinate)

        # v(t+0.5) = v(t-0.5) + a(t) * dt
        p_tilde_half = p_tilde + frc * self.time_step - self.omega_p ** 2 * x_tilde * self.time_step * self.mass_k

        # (B,A,D) = (B,A,D) - (B,1,D)
        if self.remove_com:
            p_tilde_half = p_tilde_half - self.get_pimd_com_velocity(p_tilde_half)
        kinetics = self.get_pimd_kinetics(p_tilde_half)
       
        # R(t+0.5) = R(t) + v(t+0.5) * dt
        x_tilde_half = x_tilde + p_tilde_half * 0.5 * self.time_step / self.mass_prime

        if self.temperature is not None:
            # v'(t+0.5) = f_T[v(t+0.5)]
            p_tilde_half = self.langevin_thermostat(p_tilde_half, step)

        # R(t+1) = R(t+0.5) + v'(t+0.5) * dt
        x_tilde_new = x_tilde_half + p_tilde_half * 0.5 * self.time_step / self.mass_prime

        coordinate_new = self.sqrtP*torch.einsum("jk,kad->jad",C_jk,x_tilde_new)
        momentum_half = self.sqrtP*torch.einsum("jk,kad->jad",C_jk,p_tilde_half)
        velocity_half = momentum_half / self._atom_mass

        if self.pressure is not None:
            if self.pressure_control_step == 1 or step % self.pressure_control_step == 0:
                volume = self.get_volume(pbc_box[0])
                pressure = self.get_pimd_pressure(kinetics,virial,coordinate,force) / volume 
                coordinate_new, velocity_half, pbc_box = self.andersen_barostat(coordinate_new, pressure, volume, velocity_half, pbc_box)
        return coordinate_new, velocity_half, force, energy, kinetics, virial, pbc_box
    
    def langevin_thermostat(self, momentum, step):
        if self.control_step == 1 or step % self.control_step == 0:
            momentum = self.c1 * momentum + self.c2 * torch.sqrt(1 / self.beta) * self.sqrt_mass * torch.randn(momentum.shape, device=self.device)
        return momentum
    
    def andersen_barostat(self, coordinate, pressure, volume, velocity, pbc_box):
        dv_dt = (pressure - self.pressure) / self.h_mass_inverse_0 * volume
        volume_new = volume + dv_dt * self.time_step
        crd_scale_factor = torch.pow(volume_new / volume, torch.tensor(1/3, device=volume.device))
        coordinate = coordinate * crd_scale_factor.unsqueeze(-1)
        pbc_box = pbc_box * crd_scale_factor
        velocity = velocity * torch.reciprocal(crd_scale_factor.unsqueeze(-1))
        return coordinate, velocity, pbc_box

    def get_pimd_kinetics(self, momentum):
        if momentum is None:
            return None
        # (B, A, D) * (B, A, 1)
        k = 0.5 / self.mass_prime * momentum**2
        # (B, D) <- (B, A, D)
        kinetics = torch.sum(k, -2)
        return kinetics * self.kinetic_unit_scale
    
    def get_pimd_com_velocity(self, momentum: Tensor, keepdims: bool = True) -> Tensor:

        # (B, A, D) = (B, A, D) * (B, A, 1)
        if keepdims:
            # (B, 1, D) <- (B, A, D)
            tot_vel = torch.sum(momentum, -2, keepdim=True)
            # (B, 1, 1) <- (B, 1)
            tot_mass = self.mass_prime.sum(1).unsqueeze(-1)
            com_vel = tot_vel / tot_mass * self._atom_mass
        else:
            # (B, D) <- (B, A, D)
            tot_vel = torch.sum(momentum, -2)
            # (B, 1)
            tot_mass = self.mass_prime.sum(-1)
            com_vel = tot_vel / tot_mass * self.atom_mass
        
        return com_vel
    
    def get_pimd_pressure(self,kinetics,virial,coordinate,force):
        pressure = 2 * kinetics.mean(0, keepdim=True) - 2*virial.mean(0,keepdim=True) - ((coordinate - coordinate.mean(0))*force).sum(-2).mean(0,keepdim=True)
        return pressure.mean(-1) * self.press_unit_scale